{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Zebra Puzzle\n",
    "\n",
    "This is a classic logic puzzle that is useful for demonstrating a constraint satisfaction problem (CSP).\n",
    "It is sometimes referred to as either *Einstein's Puzzle*, or *Einstein's Riddle*,\n",
    "due to an apocryphal story that it was invented by Albert Einstein as a boy.\n",
    "\n",
    "From http://en.wikipedia.org/wiki/Zebra_Puzzle:\n",
    "\n",
    "1. There are five houses.\n",
    "2. The Englishman lives in the red house.\n",
    "3. The Spaniard owns the dog.\n",
    "3. Coffee is drunk in the green house.\n",
    "4. The Ukrainian drinks tea.\n",
    "5. The green house is immediately to the right of the ivory house.\n",
    "6. The Old Gold smoker owns snails.\n",
    "7. Kools are smoked in the yellow house.\n",
    "8. Milk is drunk in the middle house.\n",
    "9. The Norwegian lives in the first house.\n",
    "10. The man who smokes Chesterfields lives in the house next to the man with the fox.\n",
    "11. Kools are smoked in the house next to the house where the horse is kept.\n",
    "12. The Lucky Strike smoker drinks orange juice.\n",
    "13. The Japanese smokes Parliaments.\n",
    "14. The Norwegian lives next to the blue house.\n",
    "\n",
    "Answer the following:\n",
    "\n",
    "1. Who drinks water?\n",
    "2. Who owns the zebra?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from pyeda.inter import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "nationalities = {\"England\", \"Japan\", \"Norway\", \"Spain\", \"Ukraine\"}\n",
    "colors = {\"blue\", \"red\", \"green\", \"ivory\", \"yellow\"}\n",
    "pets = {\"dog\", \"fox\", \"horse\", \"snails\", \"zebra\"}\n",
    "drinks = {\"coffee\", \"milk\", \"oj\", \"tea\", \"water\"}\n",
    "smokes = {\"Chesterfield\", \"OldGold\", \"Kools\", \"LuckyStrike\", \"Parliament\"}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Create Boolean variables\n",
    "xs = []\n",
    "for i in range(5):\n",
    "    xs.append({})\n",
    "    for group in (nationalities, colors, pets, drinks, smokes):\n",
    "        for name in group:\n",
    "            xs[i][name] = exprvar(name, i)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Start building a list of constraints\n",
    "cons = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Basic uniqueness constraints\n",
    "for group in (nationalities, colors, pets, drinks, smokes):\n",
    "    # Each house has exactly one {name}\n",
    "    for i in range(5):\n",
    "        cons.append(OneHot(*[xs[i][name] for name in group]))\n",
    "    # Each {name} is only in one house\n",
    "    for name in group:\n",
    "        cons.append(OneHot(*[xs[i][name] for i in range(5)]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# 2. The Englishman lives in the red house\n",
    "for i in range(5):\n",
    "    cons.append(xs[i][\"England\"] >> xs[i][\"red\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# 3. The Spaniard owns the dog\n",
    "for i in range(5):\n",
    "    cons.append(xs[i][\"Spain\"] >> xs[i][\"dog\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# 4. Coffee is drunk in the green house\n",
    "for i in range(5):\n",
    "    cons.append(xs[i][\"coffee\"] >> xs[i][\"green\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# 5. The Ukrainian drinks tea\n",
    "for i in range(5):\n",
    "    cons.append(xs[i][\"Ukraine\"] >> xs[i][\"tea\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# 6. The green house is immediately to the right of the ivory house\n",
    "cons += [\n",
    "    ~xs[0][\"green\"],\n",
    "    xs[0][\"ivory\"] >> xs[1][\"green\"],\n",
    "    xs[1][\"ivory\"] >> xs[2][\"green\"],\n",
    "    xs[2][\"ivory\"] >> xs[3][\"green\"],\n",
    "    xs[3][\"ivory\"] >> xs[4][\"green\"],\n",
    "    ~xs[4][\"ivory\"],\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# 7. The Old Gold smoker owns snails\n",
    "for i in range(5):\n",
    "    cons.append(xs[i][\"OldGold\"] >> xs[i][\"snails\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# 8. Kools are smoked in the yellow house\n",
    "for i in range(5):\n",
    "    cons.append(xs[i][\"Kools\"] >> xs[i][\"yellow\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# 9. Milk is drunk in the middle house\n",
    "cons.append(xs[2][\"milk\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# 10. The Norwegian lives in the first house\n",
    "cons.append(xs[0][\"Norway\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# 11. The man who smokes Chesterfields lives in the house next to the man with the fox\n",
    "cons += [\n",
    "    xs[0][\"Chesterfield\"] >> xs[1][\"fox\"],\n",
    "    xs[1][\"Chesterfield\"] >> (xs[0][\"fox\"] | xs[2][\"fox\"]),\n",
    "    xs[2][\"Chesterfield\"] >> (xs[1][\"fox\"] | xs[3][\"fox\"]),\n",
    "    xs[3][\"Chesterfield\"] >> (xs[2][\"fox\"] | xs[4][\"fox\"]),\n",
    "    xs[4][\"Chesterfield\"] >> xs[3][\"fox\"],\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# 12. Kools are smoked in the house next to the house where the horse is kept\n",
    "cons += [\n",
    "    xs[0][\"Kools\"] >> xs[1][\"horse\"],\n",
    "    xs[1][\"Kools\"] >> (xs[0][\"horse\"] | xs[2][\"horse\"]),\n",
    "    xs[2][\"Kools\"] >> (xs[1][\"horse\"] | xs[3][\"horse\"]),\n",
    "    xs[3][\"Kools\"] >> (xs[2][\"horse\"] | xs[4][\"horse\"]),\n",
    "    xs[4][\"Kools\"] >> xs[3][\"horse\"],\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# 13. The Lucky Strike smoker drinks orange juice\n",
    "for i in range(5):\n",
    "    cons.append(xs[i][\"LuckyStrike\"] >> xs[i][\"oj\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# 14. The Japanese smokes Parliaments\n",
    "for i in range(5):\n",
    "    cons.append(xs[i][\"Japan\"] >> xs[i][\"Parliament\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# 15. The Norwegian lives next to the blue house\n",
    "cons += [\n",
    "    xs[0][\"Norway\"] >> xs[1][\"blue\"],\n",
    "    xs[1][\"Norway\"] >> (xs[0][\"blue\"] | xs[2][\"blue\"]),\n",
    "    xs[2][\"Norway\"] >> (xs[1][\"blue\"] | xs[3][\"blue\"]),\n",
    "    xs[3][\"Norway\"] >> (xs[2][\"blue\"] | xs[4][\"blue\"]),\n",
    "    xs[4][\"Norway\"] >> xs[3][\"blue\"],\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Merge all the constraints together\n",
    "F = And(*cons).to_cnf()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# The number of clauses in the CNF\n",
    "len(F.xs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Solve the puzzle\n",
    "soln = F.satisfy_one()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Verify there is only one solution\n",
    "assert F.satisfy_count() == 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Rearrange the solution into a table\n",
    "table = [{} for i in range(5)]\n",
    "for var, val in soln.items():\n",
    "    if val:\n",
    "        if var.name in nationalities:\n",
    "            table[var.indices[0]][\"nationality\"] = var.name\n",
    "        elif var.name in colors:\n",
    "            table[var.indices[0]][\"color\"] = var.name\n",
    "        elif var.name in pets:\n",
    "            table[var.indices[0]][\"pet\"] = var.name\n",
    "        elif var.name in drinks:\n",
    "            table[var.indices[0]][\"drink\"] = var.name\n",
    "        elif var.name in smokes:\n",
    "            table[var.indices[0]][\"smoke\"] = var.name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Who drinks water?\n",
    "for i in range(5):\n",
    "    if table[i][\"drink\"] == \"water\":\n",
    "        print(\"The man from\", table[i][\"nationality\"], \"drinks water.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# Who owns the zebra?\n",
    "for i in range(5):\n",
    "    if table[i][\"pet\"] == \"zebra\":\n",
    "        print(\"The man from\", table[i][\"nationality\"], \"owns the zebra.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# The full solution\n",
    "table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.4.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
